import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from python_ai.common.xcommon import sep

sep("TEACHER'S BAD CODE")
sep('data')
m = 10
x1 = np.arange(m)
x2 = x1 * 2
x3 = x1 * 3
x4 = x1 * 4
x = np.c_[x1, x2, x3, x4]
x = pd.DataFrame(x)
print(x)

sep('label')
x_train, x_test = train_test_split(x, train_size=0.7, random_state=666)
TRAIN, TEST = 'train', 'test'
TRAIN_TEST = TRAIN + TEST
pd.options.mode.chained_assignment = None
x_train[TRAIN_TEST] = TRAIN
x_test[TRAIN_TEST] = TEST
pd.options.mode.chained_assignment = 'warn'
x = pd.concat([x_train, x_test])

sep('Usage label')
idx_train = x[TRAIN_TEST] == TRAIN
idx_test = np.invert(idx_train)
print(idx_train.sort_index())
print(idx_test.sort_index())
x_train = x[idx_train]  # later
x_test = x[idx_test]
print(x_train)
print(x_test)
